import csv
import json

def deal_nan_data_train(data):
    sum_ = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    count_ = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    mean_ = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    for i in data:
        for j in range(1, 14):
            if i[j] is not '':
                sum_[j] = sum_[j] + float(i[j])
                count_[j] += 1
    for j in range(1, 14):
        if count_[j] == 0:
            mean_[j] = 0
        else:
            mean_[j] = sum_[j] / count_[j]
    for i in data:
        for j in range(1, 14):
            if i[j] is '':
                i[j] = mean_[j]
    return data


def deal_nan_data_test(data):
    sum_ = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    count_ = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    mean_ = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    for i in data:
        for j in range(13):
            if i[j] is not '':
                sum_[j] = sum_[j] + float(i[j])
                count_[j] += 1
    for j in range(13):
        if count_[j] == 0:
            mean_[j] = 0
        else:
            mean_[j] = sum_[j] / count_[j]
    for i in data:
        for j in range(13):
            if i[j] is '':
                i[j] = mean_[j]
    return data

def generate():
    data = []
    with open('../dataset/train.txt', 'r') as f:
        for line in f.readlines():
            data_line = line.split('\t')
            if data_line != '\t':
                data_line[-1] = data_line[-1].strip()
            data.append(data_line)
    data = deal_nan_data_train(data)
    data_t = []
    with open('../dataset/test.txt', 'r') as f:
        for line in f.readlines():
            data_line = line.split('\t')
            if data_line != '\t':
                data_line[-1] = data_line[-1].strip()
            data_t.append(data_line)
    data_t = deal_nan_data_test(data_t)
    sum_vocab = []
    vocab = []
    for i in range(26):
        sum_vocab.append(1)
        vocab.append(dict())
    for i in range(len(data)):
        for j in range(26):
            if data[i][14+j] == '':
                data[i][14+j] = 'nan'
            if data[i][14+j] not in vocab[j].keys():
                vocab[j][data[i][14+j]] = sum_vocab[j]
                sum_vocab[j] += 1
            data[i][14+j] = vocab[j][data[i][14+j]]
    for i in range(len(data_t)):
        for j in range(26):
            if data_t[i][13+j] == '':
                data_t[i][13+j] = 'nan'
            if data_t[i][13+j] not in vocab[j].keys():
                vocab[j][data_t[i][13+j]] = sum_vocab[j]
                sum_vocab[j] += 1
            data_t[i][13 + j] = vocab[j][data_t[i][13 + j]]
    with open('../dataset/train_plus.csv', 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerows(data)
    with open('../dataset/test_plus.csv', 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerows(data_t)
    filename = "../dataset/numbers.json"
    with open(filename, 'w') as file_obj:
        json.dump(sum_vocab, file_obj)

generate()